import click
import os
import sys

from neuro_flow.cli.utils import wrap_async
from neuro_flow.types import LocalPath

from .root import Root


CFG_FILE = {"bash": LocalPath("~/.bashrc"), "zsh": LocalPath("~/.zshrc")}
SOURCE_CMD = {"bash": "source", "zsh": "source_zsh"}

ACTIVATION_TEMPLATE = 'eval "$(_NEURO_FLOW_COMPLETE={cmd} {exe})"'


@click.group()
def completion() -> None:
    """
    Output shell completion code.
    """


@completion.command()
@click.argument("shell", type=click.Choice(["bash", "zsh"]))
@wrap_async()
async def generate(root: Root, shell: str) -> None:
    """
    Provide instruction for shell completion generation.
    """
    root.console.print(f"Push the following line into your {CFG_FILE[shell]}")
    root.console.print(
        ACTIVATION_TEMPLATE.format(cmd=SOURCE_CMD[shell], exe=sys.argv[0])
    )


@completion.command()
@click.argument("shell", type=click.Choice(["bash", "zsh"]))
@wrap_async(pass_obj=False)
async def patch(shell: str) -> None:
    """
    Automatically patch shell configuration profile to enable completion
    """
    GENERATED_START = (
        b"# Start of generated by neuro-flow. Please do not edit this comment.\n"
    )
    GENERATED_END = (
        b"\n# End of generated by neuro-flow. Please do not edit this comment.\n"
    )

    profile_file = CFG_FILE[shell].expanduser()

    code = (
        GENERATED_START
        + os.fsencode(
            ACTIVATION_TEMPLATE.format(cmd=SOURCE_CMD[shell], exe=sys.argv[0])
        )
        + GENERATED_END
    )

    try:
        with profile_file.open("rb+") as profile:
            content = profile.read()
    except FileNotFoundError:
        content = b""

    start = content.find(GENERATED_START)
    if start != -1:
        end = content.find(GENERATED_END)
        if end == -1:
            raise click.ClickException(
                f"Malformed guarding comments. Please edit {profile_file} manually"
            )
        content = content[:start] + code + content[end + len(GENERATED_END) :]
    else:
        if content != b"":
            content += b"\n" + code
        else:
            content = code

    with profile_file.open("wb+") as profile:
        profile.write(content)
